#include <ros/ros.h>
#include <geometry_msgs/PoseStamped.h>
#include <mavros_msgs/CommandBool.h>
#include <mavros_msgs/SetMode.h>
#include <mavros_msgs/State.h>
#include <nav_msgs/Path.h>
#include <std_msgs/Bool.h>
#include <tf2/LinearMath/Quaternion.h>
#include <tf2_geometry_msgs/tf2_geometry_msgs.h>

class FormationController {
private:
    ros::NodeHandle nh_;

    // 无人机ID (0, 1, 2)
    int uav_id_;
    std::string uav_namespace_;

    // 订阅器 - 所有无人机状态
    ros::Subscriber uav0_pose_sub_, uav1_pose_sub_, uav2_pose_sub_;
    ros::Subscriber uav0_state_sub_;
    ros::Subscriber local_pose_sub_, local_state_sub_;

    // 发布器 - 无人机2和3的目标位置
    ros::Publisher setpoint_pub_;
    ros::Publisher formation_path_pub_, formation_status_pub_;

    // 服务客户端 - 用于模式切换和解锁
    ros::ServiceClient set_mode_client_;
    ros::ServiceClient arming_client_;

    // 存储当前状态
    geometry_msgs::PoseStamped uav0_pose_, uav1_pose_, uav2_pose_;
    geometry_msgs::PoseStamped local_pose_;
    mavros_msgs::State uav0_state_;
    mavros_msgs::State local_state_;

    // 编队参数
    double offset_x_, offset_y_, offset_z_;
    
    bool uav0_mission_started_, formation_activated_;
    bool uav0_airborne_, uav1_ready_, uav2_ready_;
    double takeoff_altitude_;

public:
    FormationController() :
    	uav0_mission_started_(false), formation_activated_(false),
        uav0_airborne_(false), uav1_ready_(false), uav2_ready_(false),
        takeoff_altitude_(1.0) {
        // 获取无人机ID (从命名空间获取)
        std::string ns = ros::this_node::getNamespace();
        if (ns == "/uav0") uav_id_ = 0;
        else if (ns == "/uav1") uav_id_ = 1;
        else if (ns == "/uav2") uav_id_ = 2;
        else uav_id_ = 0;
        
        uav_namespace_ = ns;
        
        ROS_INFO("Initializing formation controller for %s (ID: %d)", uav_namespace_.c_str(), uav_id_);

        // 根据无人机ID设置不同的编队偏移
        if (uav_id_ == 1) {
            offset_x_ = -2.0; offset_y_ = 0.0; offset_z_ = 0.0;
        } else if (uav_id_ == 2) {
            offset_x_ = 2.0; offset_y_ = 0.0; offset_z_ = 0.0;
        } else {
            offset_x_ = 0.0; offset_y_ = 0.0; offset_z_ = 0.0;
        }
        
        ROS_INFO("Formation offsets for %s: [%.1f, %.1f, %.1f]", 
                 uav_namespace_.c_str(), offset_x_, offset_y_, offset_z_);

        // 初始化订阅器 - 所有无人机都需要订阅所有无人机的位置
        uav0_pose_sub_ = nh_.subscribe("/uav0/mavros/local_position/pose", 10, &FormationController::uav0PoseCallback, this);
        uav1_pose_sub_ = nh_.subscribe("/uav1/mavros/local_position/pose", 10, &FormationController::uav1PoseCallback, this);
        uav2_pose_sub_ = nh_.subscribe("/uav2/mavros/local_position/pose", 10, &FormationController::uav2PoseCallback, this);
        
        // 订阅UAV0的状态（所有无人机都需要）
        uav0_state_sub_ = nh_.subscribe("/uav0/mavros/state", 10, &FormationController::uav0StateCallback, this);

        // 本地状态订阅
        local_pose_sub_ = nh_.subscribe(uav_namespace_ + "/mavros/local_position/pose", 10, &FormationController::localPoseCallback, this);
        local_state_sub_ = nh_.subscribe(uav_namespace_ + "/mavros/state", 10, &FormationController::localStateCallback, this);
        
        // 初始化发布器
        setpoint_pub_ = nh_.advertise<geometry_msgs::PoseStamped>(uav_namespace_ + "/mavros/setpoint_position/local", 10);
        formation_path_pub_ = nh_.advertise<nav_msgs::Path>(uav_namespace_ + "/formation/path", 10);
        formation_status_pub_ = nh_.advertise<std_msgs::Bool>(uav_namespace_ + "/formation/active", 10);
        
        // 初始化服务客户端
        set_mode_client_ = nh_.serviceClient<mavros_msgs::SetMode>(uav_namespace_ + "/mavros/set_mode");
        arming_client_ = nh_.serviceClient<mavros_msgs::CommandBool>(uav_namespace_ + "/mavros/cmd/arming");
        
        // 只有从无人机才需要等待和激活编队
        if (uav_id_ != 0) {
            waitForMAVROS();
        }
        
        ROS_INFO("Formation Controller for %s initialized successfully", uav_namespace_.c_str());
    }

    // 等待所有 MAVROS 节点连接
    void waitForMAVROS() {
        ROS_INFO("%s: Waiting for MAVROS connection...", uav_namespace_.c_str());
        
        int timeout = 0;
        while (ros::ok() && timeout < 100) {
            ros::spinOnce();
            if (local_state_.connected) {
                ROS_INFO("%s: MAVROS connected", uav_namespace_.c_str());
                return;
            }
            ros::Duration(0.1).sleep();
            timeout++;
        }
        
        ROS_WARN("%s: MAVROS connection timeout!", uav_namespace_.c_str());
    }

    void uav0PoseCallback(const geometry_msgs::PoseStamped::ConstPtr& msg) {
        uav0_pose_ = *msg;
        
        // 检测 UAV0 是否已经起飞
        if (!uav0_airborne_ && uav0_pose_.pose.position.z > takeoff_altitude_ - 0.5) {
            uav0_airborne_ = true;
            ROS_INFO("%s: UAV0 is airborne (altitude: %.2f)", uav_namespace_.c_str(), uav0_pose_.pose.position.z);
        }
        
        // 只有从无人机才需要计算编队
        if (uav_id_ != 0) {
            // 添加调试信息
            ROS_INFO_THROTTLE(5.0, "%s: uav0_mission_started=%d, uav0_airborne=%d, formation_activated=%d", 
                            uav_namespace_.c_str(), uav0_mission_started_, uav0_airborne_, formation_activated_);
            
            if (uav0_mission_started_ && uav0_airborne_ && !formation_activated_) {
                ROS_INFO("%s: All conditions met! Activating formation...", uav_namespace_.c_str());
                activateFormation();
            }
            
            if (formation_activated_) {
                calculateFormationSetpoint();
            }
        }
    }

    // 添加UAV0状态回调函数
    void uav0StateCallback(const mavros_msgs::State::ConstPtr& msg) {
        // 所有无人机都检测UAV0的任务状态
        if ((msg->mode == "AUTO.MISSION" || msg->mode == "AUTO.LOITER") && !uav0_mission_started_) {
            uav0_mission_started_ = true;
            ROS_INFO("%s: UAV0 mission started! (mode=%s)", uav_namespace_.c_str(), msg->mode.c_str());
        }
    }

    void uav1PoseCallback(const geometry_msgs::PoseStamped::ConstPtr& msg) {
        uav1_pose_ = *msg;
        if (!uav1_ready_ && uav_id_ == 1) {
            uav1_ready_ = true;
            ROS_INFO("%s: UAV1 position tracking started", uav_namespace_.c_str());
        }
    }
    
    void uav2PoseCallback(const geometry_msgs::PoseStamped::ConstPtr& msg) {
        uav2_pose_ = *msg;
        if (!uav2_ready_ && uav_id_ == 2) {
            uav2_ready_ = true;
            ROS_INFO("%s: UAV2 position tracking started", uav_namespace_.c_str());
        }
    }

    void localPoseCallback(const geometry_msgs::PoseStamped::ConstPtr& msg) {
        local_pose_ = *msg;
    }
    
    void localStateCallback(const mavros_msgs::State::ConstPtr& msg) {
        local_state_ = *msg;
        
        // 检测 UAV0 任务开始 (只有 UAV0 需要检测)
        if (uav_id_ == 0) {
            if ((msg->mode == "AUTO.MISSION" || msg->mode == "AUTO.LOITER") && !uav0_mission_started_) {
                uav0_mission_started_ = true;
                ROS_INFO("%s: Mission started! (mode=%s)", uav_namespace_.c_str(), msg->mode.c_str());
            }
        }
    }

    void calculateFormationSetpoint() {
        geometry_msgs::PoseStamped target_pose;
        
        // 计算目标位置 (基于 UAV0 位置 + 偏移)
        target_pose.header.stamp = ros::Time::now();
        target_pose.header.frame_id = "map";
        target_pose.pose.position.x = uav0_pose_.pose.position.x + offset_x_;
        target_pose.pose.position.y = uav0_pose_.pose.position.y + offset_y_;
        target_pose.pose.position.z = uav0_pose_.pose.position.z + offset_z_;
        target_pose.pose.orientation = uav0_pose_.pose.orientation;
        
        // 发布目标位置
        setpoint_pub_.publish(target_pose);
        
        // 发布可视化路径
        publishFormationPath(uav0_pose_, target_pose);
        
        // 发布编队状态
        std_msgs::Bool status_msg;
        status_msg.data = true;
        formation_status_pub_.publish(status_msg);
    }
    
    void activateFormation() {
        ROS_INFO("%s: Activating formation control...", uav_namespace_.c_str());
        
        // 发送初始设定值 (当前位置)
        geometry_msgs::PoseStamped init_setpoint = local_pose_;
        ROS_INFO("%s: Sending initial setpoints for 3 seconds...", uav_namespace_.c_str());
        
        ros::Time start_time = ros::Time::now();
        while (ros::Time::now() - start_time < ros::Duration(3.0) && ros::ok()) {
            setpoint_pub_.publish(init_setpoint);
            ros::Duration(0.02).sleep();
        }
        
        // 设置 OFFBOARD 模式
        if (setOffboardModeWithRetry(3)) {
            // 解锁
            ros::Duration(0.5).sleep();
            if (armDroneWithRetry(3)) {
                formation_activated_ = true;
                ROS_INFO("%s: Formation control activated! (OFFBOARD + Armed)", uav_namespace_.c_str());
            }
        }
    }
    
    bool setOffboardModeWithRetry(int max_retries) {
        for (int i = 0; i < max_retries; i++) {
            if (setOffboardMode()) {
                return true;
            }
            ros::Duration(0.5).sleep();
        }
        return false;
    }
    
    bool armDroneWithRetry(int max_retries) {
        for (int i = 0; i < max_retries; i++) {
            if (armDrone()) {
                return true;
            }
            ros::Duration(0.5).sleep();
        }
        return false;
    }
    
    bool setOffboardMode() {
        mavros_msgs::SetMode offboard_set_mode;
        offboard_set_mode.request.custom_mode = "OFFBOARD";
        
        if (set_mode_client_.call(offboard_set_mode)) {
            if (offboard_set_mode.response.mode_sent) {
                ROS_INFO("%s OFFBOARD enabled", uav_namespace_.c_str());
                return true;
            }
        }
        ROS_WARN(" Failed to set OFFBOARD", uav_namespace_.c_str());
        return false;
    }
    
    bool armDrone() {
        mavros_msgs::CommandBool arm_cmd;
        arm_cmd.request.value = true;
        
        if (arming_client_.call(arm_cmd)) {
            if (arm_cmd.response.success) {
                ROS_INFO(" Armed successfully", uav_namespace_.c_str());
                return true;
            }
        }
        ROS_WARN(" Failed to arm", uav_namespace_.c_str());
        return false;
    }
    
    void publishFormationPath(const geometry_msgs::PoseStamped& leader_pose, 
                             const geometry_msgs::PoseStamped& target_pose) {
        nav_msgs::Path path;
        path.header.stamp = ros::Time::now();
        path.header.frame_id = "map";
        
        path.poses.push_back(leader_pose);
        path.poses.push_back(target_pose);
        
        formation_path_pub_.publish(path);
    }
    
    void run() {
        ros::Rate rate(20); // 20Hz
        
        while (ros::ok()) {
            ros::spinOnce();
            
            // UAV0 不需要编队控制，只发布状态
            if (uav_id_ == 0) {
                std_msgs::Bool status_msg;
                status_msg.data = uav0_mission_started_;
                formation_status_pub_.publish(status_msg);
            }
            
            rate.sleep();
        }
    }
};

int main(int argc, char** argv) {
    ros::init(argc, argv, "formation_control_node");
    
    FormationController controller;
    controller.run();
    
    return 0;
}